#!/bin/bash
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================

set -e
BASEPATH=$(cd "$(dirname $0)/.."; pwd)
OUTPUT_PATH="${BASEPATH}/output"
BUILD_RELATIVE_PATH="build_ut"
BUILD_PATH="${BASEPATH}/${BUILD_RELATIVE_PATH}"

# print usage message
usage() {
  echo "Usage:"
  echo "sh run_test.sh [-h | --help] [-v | --verbose] [-j<N>]"
  echo "               [--ascend_install_path=<PATH>] [--ascend_3rd_lib_path=<PATH>]"
  echo "               [--platform=<PLATFORM>] [--product=<PRODUCT>]"
  echo "               [--ascendcl_c]"
  echo "               [-u | --ut] [-c | --cov]"
  echo ""
  echo "Options:"
  echo "    -h, --help     Print usage"
  echo "    -v, --verbose  Display build command"
  echo "    -j<N>          Set the number of threads used for building ACL, default is 8"
  echo "        --ascend_install_path=<PATH>"
  echo "                   Set ascend package install path, default /usr/local/Ascend/ascend-toolkit/latest"
  echo "        --ascend_3rd_lib_path=<PATH>"
  echo "                   Set ascend third_party package install path, default ./output/third_party"
  echo "        --platform=<PLATFORM>"
  echo "                   Build inference or train"
  echo "        --product=<PRODUCT>"
  echo "                   Product name"
  echo "        --ascendcl_c"
  echo "                   Run ascendcl_c test, otherwise, run normal test"
  echo "    -u, --ut       Build and execute ut"
  echo "    -c, --cov      Build ut with coverage tag"
  echo ""
}

# parse and set options
checkopts() {
  VERBOSE=""
  THREAD_NUM=8
  ASCENDCL_C="off"
  PLATFORM=""
  PRODUCT=""
  ENABLE_UT="on"
  ENABLE_COV="off"
  ENABLE_ACL_UT="off"
  ENABLE_ACL_COV="off"
  ENABLE_C_UT="off"
  ENABLE_C_COV="off"
  ENABLE_GE_COMPILE="off"
  ENABLE_ENGINES_COMPILE="off"
  ENABLE_EXECUTOR_C_COMPILE="off"
  BUILD_METADEF="off"
  ASCEND_3RD_LIB_PATH="$BASEPATH/output/third_party"
  CMAKE_BUILD_TYPE="Release"

  if [ -z "$ASCEND_INSTALL_PATH" ]; then
    ASCEND_INSTALL_PATH="/usr/local/Ascend/ascend-toolkit/latest"
  fi

  # Process the options
  parsed_args=$(getopt -a -o j:hvuc -l help,verbose,ascend_install_path:,ascend_3rd_lib_path:,platform:,product:,ascendcl_c,ut,cov -- "$@") || {
    usage
    exit 1
  }

  eval set -- "$parsed_args"

  while true; do
    case "$1" in
      -h | --help)
        usage
        exit 0
        ;;
      -j)
        THREAD_NUM="$2"
        shift 2
        ;;
      -v | --verbose)
        VERBOSE="VERBOSE=1"
        shift
        ;;
      --ascend_install_path)
        ASCEND_INSTALL_PATH="$(realpath $2)"
        shift 2
        ;;
      --ascend_3rd_lib_path)
        ASCEND_3RD_LIB_PATH="$(realpath $2)"
        shift 2
        ;;
      --platform)
        PLATFORM="$2"
        shift 2
        ;;
      --product)
        PRODUCT="$2"
        shift 2
        ;;
      --ascendcl_c)
        ASCENDCL_C="on"
        shift
        ;;
      -u | --ut)
        ENABLE_UT="on"
        shift
        ;;
      -c | --cov)
        ENABLE_UT="on"
        ENABLE_COV="on"
        shift
        ;;
      --)
        shift
        break
        ;;
      *)
        echo "Undefined option: $1"
        usage
        exit 1
        ;;
    esac
  done

  if [ "$ASCENDCL_C" = "off" ]; then
    ENABLE_EXECUTOR_C_COMPILE="off"
    ENABLE_ACL_UT="on"
    ENABLE_C_UT="off"
    if [ -z "$PLATFORM" ]; then
      PLATFORM="all"
    fi
    if [ -z "$PRODUCT" ]; then
      PRODUCT="normal"
    fi
    if [ "$ENABLE_COV" = "on" ]; then
      ENABLE_ACL_COV="on"
      ENABLE_C_COV="off"
    fi
  else
    ENABLE_EXECUTOR_C_COMPILE="on"
    ENABLE_ACL_UT="off"
    ENABLE_C_UT="on"
    if [ -z "$PLATFORM" ]; then
      PLATFORM="tiny_arm"
    fi
    if [ -z "$PRODUCT" ]; then
      PRODUCT="nano"
    fi
    if [ "$ENABLE_COV" = "on" ]; then
      ENABLE_ACL_COV="off"
      ENABLE_C_COV="on"
    fi
  fi
}

mk_dir() {
  local create_dir="$1"  # the target to make
  mkdir -pv "${create_dir}"
  echo "created ${create_dir}"
}

get_acl_type() {
  if [ "$ASCENDCL_C" = "off" ]; then
    echo "ACL"
  else
    echo "ACL_C"
  fi
}

# create build path
build_acl() {
  echo "create build directory and build $(get_acl_type)";
  mk_dir "${BUILD_PATH}"
  cd "${BUILD_PATH}"

  ENABLE_ACL_COV="$(echo ${ENABLE_ACL_COV} | tr 'a-z' 'A-Z')"
  ENABLE_ACL_UT="$(echo ${ENABLE_ACL_UT} | tr 'a-z' 'A-Z')"
  ENABLE_C_COV="$(echo ${ENABLE_C_COV} | tr 'a-z' 'A-Z')"
  ENABLE_C_UT="$(echo ${ENABLE_C_UT} | tr 'a-z' 'A-Z')"
  CMAKE_ARGS="-DBUILD_OPEN_PROJECT=True \
              -DENABLE_OPEN_SRC=True \
              -DASCENDCL_C=${ASCENDCL_C} \
              -DENABLE_GE_COMPILE=${ENABLE_GE_COMPILE} \
              -DENABLE_ENGINES_COMPILE=${ENABLE_ENGINES_COMPILE} \
              -DENABLE_EXECUTOR_C_COMPILE=${ENABLE_EXECUTOR_C_COMPILE} \
              -DBUILD_METADEF=${BUILD_METADEF} \
              -DCMAKE_BUILD_TYPE=${CMAKE_BUILD_TYPE} \
              -DCMAKE_INSTALL_PREFIX=${OUTPUT_PATH} \
              -DASCEND_INSTALL_PATH=${ASCEND_INSTALL_PATH} \
              -DASCEND_3RD_LIB_PATH=${ASCEND_3RD_LIB_PATH} \
              -DENABLE_ACL_COV=${ENABLE_ACL_COV} \
              -DENABLE_ACL_UT=${ENABLE_ACL_UT} \
              -DENABLE_C_COV=${ENABLE_C_COV} \
              -DENABLE_C_UT=${ENABLE_C_UT} \
              -DPLATFORM=${PLATFORM} \
              -DPRODUCT=${PRODUCT}"

  echo "CMAKE_ARGS=${CMAKE_ARGS}"
  cmake ${CMAKE_ARGS} ..
  if [ $? -ne 0 ]; then
    echo "execute command: cmake ${CMAKE_ARGS} .. failed."
    return 1
  fi

  if [ "$ASCENDCL_C" = "off" ]; then
    make ${VERBOSE} ascendcl_utest -j${THREAD_NUM}
  else
    make ${VERBOSE} ascendcl_c_utest -j${THREAD_NUM}
  fi
  if [ $? -ne 0 ]; then
    echo "execute command: make ${VERBOSE} -j${THREAD_NUM} failed."
    return 1
  fi
  echo "$(get_acl_type) build success!"
}

run_ut_acl() {
  if [[ "X$ENABLE_UT" = "Xon" || "X$ENABLE_COV" = "Xon" ]]; then
    cp ${BUILD_PATH}/tests/ut/acl/ascendcl_utest ${OUTPUT_PATH}

    local report_dir="${OUTPUT_PATH}/report/ut" && mk_dir "${report_dir}"
    RUN_TEST_CASE="${OUTPUT_PATH}/ascendcl_utest --gtest_output=xml:${report_dir}/ascendcl_utest.xml" && ${RUN_TEST_CASE}
    if [[ "$?" -ne 0 ]]; then
      echo "!!! UT FAILED, PLEASE CHECK YOUR CHANGES !!!"
      echo -e "\033[31m${RUN_TEST_CASE}\033[0m"
      exit 1;
    fi
    echo "Generated coverage statistics, please wait..."
    cd ${BASEPATH}
    rm -rf ${BASEPATH}/cov
    mkdir ${BASEPATH}/cov
    lcov -c -d ${BUILD_RELATIVE_PATH}/tests/ut/acl -o cov/tmp.info
    lcov -r cov/tmp.info '*/output/*' "*/${BUILD_RELATIVE_PATH}/opensrc/*" "*/${BUILD_RELATIVE_PATH}/proto/*" \
        '*/third_party/*' '*/tests/*' '/usr/local/*' '/usr/include/*' \
        "${ASCEND_INSTALL_PATH}/*" "${ASCEND_3RD_LIB_PATH}/*" -o cov/coverage.info
    cd ${BASEPATH}/cov
    genhtml coverage.info
  fi
}

run_ut_acl_c() {
  if [[ "X$ENABLE_UT" = "Xon" || "X$ENABLE_COV" = "Xon" ]]; then
    cp -rf ${BUILD_PATH}/tests/ut/acl/testcase_c/ascendcl_c_utest ${OUTPUT_PATH}

    local report_dir="${OUTPUT_PATH}/report/ut" && mk_dir "${report_dir}"
    RUN_TEST_CASE="${OUTPUT_PATH}/ascendcl_c_utest --gtest_output=xml:${report_dir}/ascendcl_c_utest.xml" && ${RUN_TEST_CASE}
    if [[ "$?" -ne 0 ]]; then
      echo "!!! UT FAILED, PLEASE CHECK YOUR CHANGES !!!"
      echo -e "\033[31m${RUN_TEST_CASE}\033[0m"
      exit 1;
    fi
    echo "Generated coverage statistics, please wait..."
    cd ${BASEPATH}
    rm -rf ${BASEPATH}/cov
    mkdir ${BASEPATH}/cov
    lcov -c -d ${BUILD_RELATIVE_PATH}/tests/ut/acl/testcase_c -o cov/tmp.info
    lcov -r cov/tmp.info '*/output/*' "*/${BUILD_RELATIVE_PATH}/opensrc/*" "*/${BUILD_RELATIVE_PATH}/proto/*" \
        '*/third_party/*' '*/tests/*' '/usr/local/*' '/usr/include/*' \
        "${ASCEND_INSTALL_PATH}/*" "${ASCEND_3RD_LIB_PATH}/*" -o cov/coverage.info
    cd ${BASEPATH}/cov
    genhtml coverage.info
  fi
}

main() {
  checkopts "$@"

  # Acl build start
  echo "---------------- $(get_acl_type) build start ----------------"
  g++ -v
  mk_dir ${OUTPUT_PATH}
  build_acl
  if [[ "$?" -ne 0 ]]; then
    echo "$(get_acl_type) build failed.";
    exit 1;
  fi
  echo "---------------- $(get_acl_type) build finished ----------------"

  rm -f ${OUTPUT_PATH}/libgmock*.so
  rm -f ${OUTPUT_PATH}/libgtest*.so
  rm -f ${OUTPUT_PATH}/lib*_stub.so

  chmod -R 750 ${OUTPUT_PATH}

  echo "---------------- $(get_acl_type) output generated ----------------"
  if [ "$ASCENDCL_C" = "off" ]; then
    run_ut_acl
  else
    run_ut_acl_c
  fi
}

main "$@"
